import argparse
from typing import List, Dict, Optional, Union, Any
from collections import defaultdict


def parse_args(arg_names: List[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    for arg_name in arg_names:
        parser.add_argument("--" + arg_name)
    args = parser.parse_args()
    return args


def optional_type(_type):
    def _type_func(val_str: Optional[str]):
        if val_str is None:
            return None
        val_str = val_str.strip()
        if val_str in ["null", "None"]:
            return None
        return _type(val_str)

    return _type_func


class ArgDef:
    def __init__(self, _name, _type, _default):
        """
        example:
        ArgDef("batch_size",int,64)
        """
        self.name = _name
        self.type = _type
        self.default = _default


class Args:
    def __init__(self, arg_def_list: List[ArgDef]):
        """
        args_def example:
        [
            ArgDef("batch_size",int,64),
            ArgDef("learning_rate",float,1e-3)
        ]
        """

        _args_def = {}
        for arg_def in arg_def_list:
            _args_def[arg_def.name] = {
                "type": arg_def.type,
                "default": arg_def.default
            }
        _arg_names = list(_args_def)
        _raw_args = parse_args(_arg_names)
        _args = defaultdict()
        for arg_name in _raw_args.__dict__:
            info = _args_def.get(arg_name)
            _type = info["type"]
            _default = info["default"]
            if _type is None:
                setattr(self, arg_name, _raw_args.__dict__[arg_name])
            else:
                raw_val = getattr(_raw_args, arg_name, None)
                if raw_val is None:
                    setattr(self, arg_name, _default)
                else:
                    if _type is bool:
                        v = False
                        if raw_val in ["false", "False"]:
                            v = False
                        elif raw_val in ["true", "True"]:
                            v = True
                        else:
                            v = bool(int(raw_val))
                        setattr(self, arg_name, v)
                    else:
                        setattr(self, arg_name, _type(raw_val))

    @property
    def args(self):
        return self.__dict__


D = ArgDef

if __name__ == "__main__":
    args = Args([
        D("batch_size", int, 64),
        D("learning_rate", float, 1e-3),
        D("num_epochs", int, 10),
        D("data_dir", str, None),
        D("save_dir", str, None)
    ])
    print(args.args)
    print(args.batch_size)
